#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 内存Cookie提取工具主程序

import os
import sys
import argparse
import logging
from datetime import datetime

# 导入自定义模块
try:
    from process_finder import ProcessFinder
    from cookie_extractor import CookieExtractor
    from enhanced_pattern_scanner import PatternScanner
except ImportError:
    print("错误: 无法导入必要的模块。请确保所有文件都在同一目录中。")
    sys.exit(1)

# 设置日志
def setup_logging(verbose=False):
    level = logging.INFO if verbose else logging.WARNING
    logging.basicConfig(
        level=level,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        handlers=[
            logging.StreamHandler(),
            logging.FileHandler("cookie_extractor.log")
        ]
    )
    return logging.getLogger("cookie_extractor")

def main():
    parser = argparse.ArgumentParser(description='内存Cookie提取工具')
    subparsers = parser.add_subparsers(dest='command', help='命令')
    
    # 查找进程子命令
    find_parser = subparsers.add_parser('find', help='查找进程')
    find_parser.add_argument('-l', '--list', action='store_true', help='列出所有进程')
    find_parser.add_argument('-n', '--name', type=str, help='按名称查找进程')
    find_parser.add_argument('-c', '--cookie', action='store_true', help='查找可能包含Cookie的进程')
    find_parser.add_argument('-p', '--pid', type=int, help='获取特定PID的进程详情')
    find_parser.add_argument('-d', '--detailed', action='store_true', help='显示详细信息')
    
    # 提取Cookie子命令
    extract_parser = subparsers.add_parser('extract', help='从进程中提取Cookie')
    extract_parser.add_argument('target', help='目标进程名称或PID')
    extract_parser.add_argument('-o', '--output-dir', default='cookies_extracted', help='输出目录')
    extract_parser.add_argument('-a', '--auto-pattern', action='store_true', help='使用增强模式扫描')
    
    # 全局选项
    parser.add_argument('-v', '--verbose', action='store_true', help='显示详细日志')
    
    args = parser.parse_args()
    
    # 设置日志
    logger = setup_logging(args.verbose)
    
    # 如果没有指定命令，显示帮助
    if not args.command:
        parser.print_help()
        return
    
    # 处理查找进程命令
    if args.command == 'find':
        logger.info("启动进程查找器")
        finder = ProcessFinder()
        
        if args.list:
            processes = finder.list_all_processes()
            print(f"发现 {len(processes)} 个进程:")
            for proc in processes:
                finder.print_process_info(proc, args.detailed)
        
        elif args.name:
            processes = finder.find_processes_by_name(args.name)
            print(f"名称包含 '{args.name}' 的进程 ({len(processes)} 个):")
            for proc in processes:
                finder.print_process_info(proc, args.detailed)
        
        elif args.cookie:
            processes = finder.find_likely_cookie_processes()
            print(f"可能包含Cookie的进程 ({len(processes)} 个):")
            for proc in processes:
                finder.print_process_info(proc, args.detailed)
        
        elif args.pid:
            process = finder.get_process_details(args.pid)
            print(f"PID {args.pid} 的进程详情:")
            finder.print_process_info(process, True)
        
        else:
            find_parser.print_help()
    
    # 处理提取Cookie命令
    elif args.command == 'extract':
        logger.info(f"开始从目标 '{args.target}' 提取Cookie")
        
        # 解析目标进程
        target = args.target
        try:
            target = int(target)
            logger.info(f"目标指定为PID: {target}")
        except ValueError:
            logger.info(f"目标指定为进程名: {target}")
        
        # 创建输出目录
        os.makedirs(args.output_dir, exist_ok=True)
        
        # 创建提取器
        if args.auto_pattern:
            logger.info("使用增强模式扫描")
            
            # 先使用常规提取器获取内存字符串
            print("第1阶段: 使用Frida提取内存...")
            extractor = CookieExtractor(output_dir=args.output_dir)
            
            try:
                # 执行提取
                extractor.extract_from_process(target)
                
                # 第二阶段: 使用增强模式进一步分析
                print("\n第2阶段: 使用增强模式分析...")
                
                # 分析每个提取的cookie文件
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                scanner = PatternScanner()
                
                for user_id, cookies in extractor.cookies_by_user.items():
                    # 合并所有cookie字符串用于深度分析
                    cookie_text = "\n".join(cookies)
                    results = scanner.scan_text(cookie_text)
                    scanner.add_results(user_id, results)
                
                # 保存增强分析结果
                scanner.save_results(args.output_dir)
                
            except Exception as e:
                logger.error(f"提取过程中出错: {e}")
                print(f"提取过程中出错: {e}")
                
        else:
            logger.info("使用标准模式提取")
            extractor = CookieExtractor(output_dir=args.output_dir)
            
            try:
                # 执行提取
                extractor.extract_from_process(target)
            except Exception as e:
                logger.error(f"提取过程中出错: {e}")
                print(f"提取过程中出错: {e}")

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n操作已取消")
    except Exception as e:
        print(f"发生错误: {e}")
        if '--verbose' in sys.argv or '-v' in sys.argv:
            import traceback
            traceback.print_exc() 